#include <algorithm>
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>

#include "opencv2/opencv.hpp"

#include "Eigen/Core"
#include "Eigen/Dense"

#include "glog/logging.h"

// chessboard
const int cb_pattern_height = 7;
const int cb_pattern_width = 5;
const double cb_square_len = 95;
const double board_dim_height = 850;
const double board_dim_width = 610;
const double board_trans_error_x = 2;
const double board_trans_error_y = 0;

auto ChessBoardProjection(const std::vector<cv::Point2d>& corners,
                          const cv::Mat& img) {
  // find the chessboard in 3D space - in it's own object frame
  // (position is arbitrary, so we place in flat)

  // location of boatd frame origin from the bottom right inner corner of the chessboard
  // board frame: origin at center, x-axis: left-to-right, y-axis: top-to-bottom
  // board z-axis is always zero
  cv::Point3d cb_br_corner((cb_pattern_width - 1) * cb_square_len / 2,
                           (cb_pattern_height - 1) * cb_square_len / 2, 0);
  // all chessboard inner corners in 3D
  std::vector<cv::Point3d> cb_corners_3d;
  for (int y = 0; y < cb_pattern_height; y++) {
    for (int x = 0; x < cb_pattern_width; x++) {
      // first point in top-left
      cb_corners_3d.push_back(
          cv::Point3d(cv::Point3d(x, y, 0) * cb_square_len - cb_br_corner));
    }
  }

  // chessboard corners, middle square corners, board corners and center
  std::vector<cv::Point3d> board_corners_3d;
  // board corner coordinates from the ceter of the chessboard
  // board right-bottom corner
  board_corners_3d.push_back(
      cv::Point3d((board_dim_width - board_trans_error_x) / 2.0,
                  (board_dim_height - board_trans_error_y) / 2.0, 0.0));
  // board left-bottom corner
  board_corners_3d.push_back(
      cv::Point3d(-(board_dim_width + board_trans_error_x) / 2.0,
                  (board_dim_height - board_trans_error_y) / 2.0, 0.0));
  // board left-top corner
  board_corners_3d.push_back(
      cv::Point3d(-(board_dim_width + board_trans_error_x) / 2.0,
                  -(board_dim_height + board_trans_error_y) / 2.0, 0.0));
  // board right-top corner
  board_corners_3d.push_back(
      cv::Point3d((board_dim_width - board_trans_error_x) / 2.0,
                  -(board_dim_height + board_trans_error_y) / 2.0, 0.0));
  // board center
  board_corners_3d.push_back(
      cv::Point3d(-board_trans_error_x / 2.0, -board_trans_error_y / 2.0, 0.0));

  // chessboard and board corners in image
  std::vector<cv::Point2d> cb_inner_corner_pixels, board_corner_pixels;
  // fisheye model
  // undistort the image by applying the fisheye intrinsic parameters
  // the final input params is the camera matrix in the new or rectified coordinate frame.
  // we put this to be the same as loaded camera intrinsic parameters or else it will be
  // set to empty matrix by default.
  cv::Mat rvec(3, 3, cv::DataType<double>::type);
  cv::Mat tvec(3, 1, cv::DataType<double>::type);

  cv::Mat camera_mat, dist_coeff;
  camera_mat = (cv::Mat_<double>(3, 3) << 1176.93, 0.0, 962.275, 0.0, 1177.72,
                612.725, 0.0, 0.0, 1);
  dist_coeff =
      (cv::Mat_<double>(1, 4) << -0.0540096, -0.0784275, 0.0959641, -0.0515253);

  std::vector<cv::Point2d> corners_undistored;
  cv::fisheye::undistortPoints(corners, corners_undistored, camera_mat,
                               dist_coeff, cv::noArray(), camera_mat);
  // chessboard inner corners(3D) - undistored 2D point in image
  cv::solvePnP(cb_corners_3d, corners_undistored, camera_mat, cv::noArray(),
               rvec, tvec);
  // chessboard inner corners projects to image
  cv::fisheye::projectPoints(cb_corners_3d, cb_inner_corner_pixels, rvec, tvec,
                             camera_mat, dist_coeff);
  // board corners projects to image
  cv::fisheye::projectPoints(board_corners_3d, board_corner_pixels, rvec, tvec,
                             camera_mat, dist_coeff);

  // display projected corners
  cv::Mat img_show = img.clone();
  // projected board corners in image
  for (int i = 0; i < board_corner_pixels.size(); i++) {
    if (i == 0) {
      cv::circle(img_show, board_corner_pixels[i], 4, CV_RGB(255, 0, 0), -1);
    } else if (i == 1) {
      cv::circle(img_show, board_corner_pixels[i], 4, CV_RGB(0, 255, 0), -1);
    } else if (i == 2) {
      cv::circle(img_show, board_corner_pixels[i], 4, CV_RGB(0, 0, 255), -1);
    } else if (i == 3) {
      cv::circle(img_show, board_corner_pixels[i], 4, CV_RGB(255, 255, 0), -1);
    } else if (i == 4) {
      cv::circle(img_show, board_corner_pixels[i], 4, CV_RGB(0, 255, 255), -1);
    }
  }
  // projected chessboard corners in image
  for (auto& point : cb_inner_corner_pixels) {
    cv::circle(img_show, point, 3, CV_RGB(255, 0, 0), -1);
  }

  // debug
  cv::imshow("projected corners", img_show);
  cv::waitKey(0);

  // return all the necessary coefficients
  return std::make_tuple(rvec, tvec, board_corners_3d);
}

// extract chessboard
std::tuple<std::vector<cv::Point3d>, cv::Mat> LocateChessBoard(
    const cv::Mat& img) {
  cv::Mat gray_img;
  cv::cvtColor(img, gray_img, cv::COLOR_BGR2GRAY);
  std::vector<cv::Point2f> cornersf;
  std::vector<cv::Point2d> corners;
  // find chessboard pattern in the image
  bool pattern_found = cv::findChessboardCorners(
      gray_img, cv::Size(cb_pattern_width, cb_pattern_height), cornersf,
      cv::CALIB_CB_ADAPTIVE_THRESH + cv::CALIB_CB_NORMALIZE_IMAGE);
  if (!pattern_found) {
    LOG(INFO) << "No chessboard found";
    std::vector<cv::Point3d> empty_corners;
    cv::Mat empty_normals;
    return std::make_tuple(empty_corners, empty_normals);
  }
  LOG(INFO) << "chessboard found";

  // find corner points with sub-pixel accuracy
  // this throws a exception if the corner points are doubles and not floats
  cv::cornerSubPix(
      gray_img, cornersf, cv::Size(11, 11), cv::Size(-1, -1),
      cv::TermCriteria(cv::TermCriteria::MAX_ITER | cv::TermCriteria::EPS, 30,
                       0.1));

  for (const auto& c : cornersf) {
    corners.push_back(cv::Point2d(c));
  }

  auto [rvec, tvec, board_corners_3d] = ChessBoardProjection(corners, img);

  cv::Mat rmat;
  cv::Rodrigues(rvec, rmat);
  // TODO
  cv::Mat z = cv::Mat(cv::Point3d(0.0, 0.0, -1.0));
  auto chessboard_normal = rmat * z;

  // TODO
  std::vector<cv::Point3d> corner_vectors;
  for (const auto& c : board_corners_3d) {
    cv::Mat m(rmat * cv::Mat(c).reshape(1) + tvec);
    corner_vectors.push_back(cv::Point3d(m));
  }

  return std::make_tuple(corner_vectors, chessboard_normal);
}

int main(int argc, char** argv) {
  // load img
  cv::Mat img = cv::imread(argv[1]);
  LocateChessBoard(img);
  return 0;
}
